import os
from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv
import requests
import json
import streamlit as st

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPUAI_API_KEY')
amap_key = os.environ.get('amap_key')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")
client = ZhipuAI(api_key=api_key)

def get_completion(messages, model="glm-4"):
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.01,  
        tool_choice="auto", 
        tools=[{
            "type": "web_search",
            "web_search": {
                "enable": True,
                "search_result": True,
            }
        }],
    )
    print(f"\nget_completion:{response}")
    return response


# 头像配置
ICON_AI = '💻'
ICON_USER = '🧑'

# 显示一条消息（包含头像与消息内容）
def dspMessage(role, content):
    with st.chat_message(role, avatar=ICON_AI if role == 'assistant' else ICON_USER):
        st.write(content)
        

# 追加并显示一条消息
def append_and_show(role, content):
    st.session_state.messages.append({"role": role, "content": content})      
    dspMessage(role, content)

if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": """assistant""", "content": """
        我是你的信息检索小助手，请问你要搜索什么信息？"""}]


# 将会话中的messages列表中的消息全部显示出来
for msg in st.session_state.messages:
    dspMessage(msg["role"], msg["content"])

# 接受用户输入的提示词，并调用大模型API获得反馈
if prompt := st.chat_input():    
    append_and_show("user", prompt)

    messages = [
        {"role": "system", 
          "content": """你是一个具备网络访问能力的信息小助手，
          你应当先分析用户消息中的关键主题，然后优先使用网检索相关信息，
          然后根据搜索的结果给出回答。"""},
        {"role": "user", "content": prompt}
    ]

    response = get_completion(messages)

    content = response.choices[0].message.content
    print(f"\n综合结果：{content}\n")
    if (hasattr(response, "web_search") and 
        response.web_search is not None):
        content = f"{content}\n\n*以下是消息来源：*\n"

        # 取得检索到的网页信息
        for searchResult in response.web_search:
            print(f"\nsearchResult=\n{searchResult}\n")
            webContent = searchResult["content"]
            webIcon = None
            if "icon"in searchResult:
                webIcon = searchResult["icon"]
            link = searchResult["link"]
            media = None
            if "media" in searchResult:
                media = searchResult["media"]
            title = searchResult["title"]

            print(f"\nwebConten={webContent}\nwebIcon={webIcon}\nlink={link}\nmedia={media}\ntitle={title}")
            if media is None:
                searchContent = f"[{title}]({link})"
            else:
                searchContent = f"[{title} - {media}]({link})"
            content = f"{content}\n{searchContent}\n"


    # 保存并显示已经完成的回复
    append_and_show("assistant", content) 
    print(f"\n===END===\n")

